from fastapi import APIRouter, WebSocket, WebSocketDisconnect
import json
import asyncio
import random
from datetime import datetime, timedelta
import time
from websockets.exceptions import ConnectionClosedOK, ConnectionClosedError

router = APIRouter(tags=["websocket"])

# Store connected clients
connected_clients = []

# Store time series data points (limited to most recent points only)
historical_time_series = []
MAX_HISTORY_POINTS = 12  # Store only recent data points for chart display

# Mock data for RPA tasks
bank_names = ["工商银行", "建设银行", "农业银行", "中国银行", "交通银行", "招商银行", "浦发银行", "民生银行", "兴业银行", "平安银行","中原银行", "华润银行", "湖北银行", "广发银行", "东莞银行", "广州银行"]
status_options = ["成功", "失败", "运行中"]

# Generate mock data
def generate_mock_data():
    global historical_time_series
    now = datetime.now()
    
    # First, decide how many tasks of each status to generate
    # Make sure the total equals the number of banks
    num_banks = len(bank_names)
    current_running = random.randint(2, 5)
    current_failed = random.randint(1, 3)
    current_success = num_banks - current_running - current_failed
    
    # Ensure success is always positive
    if current_success < 0:
        # Adjust if needed
        overflow = abs(current_success)
        current_running = max(1, current_running - overflow//2)
        current_failed = max(1, current_failed - (overflow - overflow//2))
        current_success = num_banks - current_running - current_failed
    
    # Generate bank tasks data
    bank_tasks = []
    
    # Distribute statuses randomly among banks
    all_statuses = ["成功"] * current_success + ["失败"] * current_failed + ["运行中"] * current_running
    random.shuffle(all_statuses)
    
    # Generate tasks for each bank with the pre-determined statuses
    for i, bank in enumerate(bank_names):
        status = all_statuses[i] if i < len(all_statuses) else "成功"  # Default to success if we run out
        start_time = now - timedelta(minutes=random.randint(5, 60))
        
        task = {
            "bank_name": bank,
            "status": status,
            "start_time": start_time.strftime("%Y-%m-%d %H:%M:%S"),
        }
        
        if status != "运行中":
            end_time = start_time + timedelta(minutes=random.randint(1, 10))
            duration = (end_time - start_time).total_seconds()
            task["end_time"] = end_time.strftime("%Y-%m-%d %H:%M:%S")
            task["duration"] = f"{int(duration // 60)}分{int(duration % 60)}秒"
        else:
            task["end_time"] = "-"
            task["duration"] = "-"
            
        if status == "失败":
            errors = ["网络连接超时", "数据格式错误", "认证失败", "服务器无响应", "数据不一致"]
            task["error_reason"] = random.choice(errors)
        else:
            task["error_reason"] = "-"
            
        bank_tasks.append(task)
    
    # Count the final number of each status (to handle any edge cases)
    actual_success = sum(1 for task in bank_tasks if task["status"] == "成功")
    actual_running = sum(1 for task in bank_tasks if task["status"] == "运行中")
    actual_failed = sum(1 for task in bank_tasks if task["status"] == "失败")
    
    # Create a new time series data point with the actual counts
    new_data_point = {
        "time": now.strftime("%H:%M:%S"),
        "success": actual_success,
        "running": actual_running,
        "failed": actual_failed
    }
    
    # Add to historical data
    historical_time_series.append(new_data_point)
    
    # Trim historical data to keep only the most recent MAX_HISTORY_POINTS
    if len(historical_time_series) > MAX_HISTORY_POINTS:
        historical_time_series = historical_time_series[-MAX_HISTORY_POINTS:]
    
    return {
        "summary": {
            "success": actual_success,
            "running": actual_running,
            "failed": actual_failed
        },
        "time_series": historical_time_series,
        "bank_tasks": bank_tasks,
        "timestamp": now.strftime("%Y-%m-%d %H:%M:%S")
    }

@router.websocket("/ws")
async def websocket_endpoint(websocket: WebSocket):
    await websocket.accept()
    connected_clients.append(websocket)
    
    try:
        while True:
            # Send data every 5 seconds
            data = generate_mock_data()
            print("data",json.dumps(data,ensure_ascii=False))
            await websocket.send_json(data)
            await asyncio.sleep(5)
    except (WebSocketDisconnect, ConnectionClosedOK, ConnectionClosedError):
        # Handle normal client disconnection
        if websocket in connected_clients:
            connected_clients.remove(websocket)
    except Exception as e:
        # Handle other exceptions
        print(f"WebSocket error: {str(e)}")
        if websocket in connected_clients:
            connected_clients.remove(websocket)

# Broadcast to all connected clients
async def broadcast_data():
    while True:
        if connected_clients:
            data = generate_mock_data()
            # Create a copy of the list to avoid modification during iteration
            clients = connected_clients.copy()
            for client in clients:
                try:
                    await client.send_json(data)
                except (WebSocketDisconnect, ConnectionClosedOK, ConnectionClosedError):
                    # Client disconnected, remove it from the list
                    if client in connected_clients:
                        connected_clients.remove(client)
                except Exception as e:
                    print(f"Error broadcasting to client: {str(e)}")
                    if client in connected_clients:
                        connected_clients.remove(client)
        await asyncio.sleep(5) 